import copy
import numpy as np
import torch
import time
from flcore.clients.clientcp import *
from utils.data_utils import partition_indices, get_subset_size
from utils.data_utils import partition_indices, _load_full
# ---------- FedCP.__init__ 开头 ----------
from flcore.clients.fedcp_wrapper import wrap
from threading import Thread


class FedCP:
    def __init__(self, args, times):#times表实验次数或标识符
        self.device = args.device#args包含设备，数据集，全局轮次，模型
        self.dataset = args.dataset
        self.global_rounds = args.global_rounds
        # 1) 保证 model 已实例化、带真实构造参数
        if isinstance(args.model, type):  # 传进来的是类
            model_kwargs = getattr(args, "model_kwargs", {})  # 可在主脚本里写入
            args.model = args.model(args, **model_kwargs)

        # 2) 若模型还没拆 head / feature_extractor，再自动 wrap
        if not (hasattr(args.model, "head") and hasattr(args.model, "feature_extractor")):
            args.model = wrap(args.model)

        # 3) 全局只维护共享的 feature_extractor 即可
        self.global_modules = copy.deepcopy(args.model.feature_extractor)

        self.num_clients = args.num_clients#总客户端数
        # -------- 一次性划分索引 --------
        train_parts = partition_indices("train", self.num_clients,
                                        strategy="iid", seed=args.seed)
        test_parts = partition_indices("test", self.num_clients,
                                       strategy="iid", seed=args.seed)

        self.join_ratio = args.join_ratio#参与比例
        self.random_join_ratio = args.random_join_ratio
        self.join_clients = int(self.num_clients * self.join_ratio)#记录每轮参与的客户端数

        self.clients = []#存储所有客户端
        self.selected_clients = []#每轮选中的客户端

        self.uploaded_weights = []
        self.uploaded_ids = []
        self.uploaded_models = []

        self.rs_test_acc = []#记录测试准确率
        self.rs_train_loss = []

        self.times = times
        self.eval_gap = args.eval_gap

        in_dim = list(args.model.head.parameters())[0].shape[1]
        cs = ConditionalSelection(in_dim, in_dim).to(args.device)#创建条件选择模块

        #创建客户端
        for i in range(self.num_clients):
            client = clientCP(args, 
                            id=i,
                            train_idx=train_parts[i],
                            test_idx=test_parts[i],
                            ConditionalSelection=cs)#导入条件选择模块
            self.clients.append(client)

        print(f"\nJoin ratio / total clients: {self.join_ratio} / {self.num_clients}")
        print("Finished creating server and clients.")

        # self.load_model()
        self.Budget = []#记录每轮耗时
        self.head = None
        self.cs = None

    #选择客户端（动态选择
    def select_clients(self):#动态选择参与方
        #join_clients是指参与测试的客户端数量
        if self.random_join_ratio:
            join_clients = np.random.choice(range(self.join_clients, self.num_clients+1), 1, replace=False)[0]
        else:
            #否则使用固定的随机选择数量
            join_clients = self.join_clients
        selected_clients = list(np.random.choice(self.clients, join_clients, replace=False))

        return selected_clients

    def send_models(self):#将全局模型发送给所有客户端
        assert (len(self.clients) > 0)

        for client in self.clients:
            #设置client的参数
            #注意是遍历所有客户端，并set_parameters
            client.set_parameters(self.global_modules)
    #将客户端模型参数按照权重w加到全局模型参数上
    def add_parameters(self, w, client_model):
        for server_param, client_param in zip(self.global_modules.parameters(), client_model.parameters()):
            server_param.data += client_param.data.clone() * w
    #初始化全局模型为第一个上传模型的部分，并清空参数，即保留参数结构
    def aggregate_parameters(self):#聚合全局模型
        assert (len(self.uploaded_models) > 0)
        #此时全局模型复制参数
        self.global_modules = copy.deepcopy(self.uploaded_models[0])
        #依照参数模式改为zero_like的参数
        for param in self.global_modules.parameters():
            param.data = torch.zeros_like(param.data)
        #w是权重，使用更新模型更新全局参数
        for w, client_model in zip(self.uploaded_weights, self.uploaded_models):
            self.add_parameters(w, client_model)
    #从每个客户端获取正确的预测数、样本数和auc
    def test_metrics(self):
        num_samples = []
        tot_correct = []
        tot_auc = []
        for c in self.clients:
            ct, ns, auc = c.test_metrics()#由客户端实现的存储测试结果的矩阵
            print(f'Client {c.id}: Acc: {ct*1.0/ns}, AUC: {auc}')
            tot_correct.append(ct*1.0)
            tot_auc.append(auc*ns)
            num_samples.append(ns)

        ids = [c.id for c in self.clients]

        return ids, num_samples, tot_correct, tot_auc
    #计算所有客户端加权平均的准确率和auc，记录结果
    def evaluate(self, acc=None):
        stats = self.test_metrics()

        test_acc = sum(stats[2])*1.0 / sum(stats[1])
        test_auc = sum(stats[3])*1.0 / sum(stats[1])
        
        if acc == None:
            self.rs_test_acc.append(test_acc)
        else:
            acc.append(test_acc)

        print("Averaged Test Accurancy: {:.4f}".format(test_acc))
        print("Averaged Test AUC: {:.4f}".format(test_auc))

    #训练循环
    def train(self):
        for i in range(self.global_rounds+1):
            #记录每轮训练用时
            s_t = time.time()
            print(f"\n▶ Global round {i}/{self.global_rounds}")  # ← 轮次提示
            #选择客户端
            self.selected_clients = self.select_clients()
            #每隔eval_gap轮评估模型
            if i%self.eval_gap == 0:
                print(f"\n-------------Round number: {i}-------------")
                print("\nEvaluate before local training")
                self.evaluate()

            for client in self.selected_clients:
                #训练本地条件选择模块
                print(f"  • Train on Client {client.id}")
                client.train_cs_model()
                #生成待上传的头部模型
                client.generate_upload_head()
            #接收客户端上传的特征提取器
            self.receive_models()
            #聚合参数
            self.aggregate_parameters()
            #向各个客户端发送模型
            self.send_models()
            #聚合并更新选中客户端的head_g和条件选择模块
            self.global_head()
            self.global_cs()
            #记录耗时
            self.Budget.append(time.time() - s_t)
            print('-'*50, self.Budget[-1])

        print("\nBest accuracy.")
        print(max(self.rs_test_acc))
        print(sum(self.Budget[1:])/len(self.Budget[1:]))

    #模型接受和更新
    def receive_models(self):
        assert (len(self.selected_clients) > 0)

        active_train_samples = 0
        for client in self.selected_clients:
            #计算训练的样本总量
            active_train_samples += client.train_samples
        #保存更新权重
        self.uploaded_weights = []
        self.uploaded_ids = []
        self.uploaded_models = []
        for client in self.selected_clients:
            #计算更新权重
            self.uploaded_weights.append(client.train_samples / active_train_samples)
            #记录用于更新的客户端id
            self.uploaded_ids.append(client.id)
            #记录特征提取器
            self.uploaded_models.append(client.model.model.feature_extractor)
    #更新选中客户端的头
    def global_head(self):
        self.uploaded_model_gs = []
        for client in self.selected_clients:
            self.uploaded_model_gs.append(client.model.head_g)
        #初始化全局头的形状
        self.head = copy.deepcopy(self.uploaded_model_gs[0])
        for param in self.head.parameters():
            param.data = torch.zeros_like(param.data)
        #将全局头作聚合
        for w, client_model in zip(self.uploaded_weights, self.uploaded_model_gs):
            self.add_head(w, client_model)
        #将客户端的头作聚合
        for client in self.selected_clients:
            client.set_head_g(self.head)

    def add_head(self, w, head):
        for server_param, client_param in zip(self.head.parameters(), head.parameters()):
            server_param.data += client_param.data.clone() * w
    #同样的方式更新全局的条件选择模块，并下放到本地
    def global_cs(self):
        self.uploaded_model_gs = []
        for client in self.selected_clients:
            self.uploaded_model_gs.append(client.model.gate.cs)

        self.cs = copy.deepcopy(self.uploaded_model_gs[0])
        for param in self.cs.parameters():
            param.data = torch.zeros_like(param.data)
            
        for w, client_model in zip(self.uploaded_weights, self.uploaded_model_gs):
            self.add_cs(w, client_model)

        for client in self.selected_clients:
            client.set_cs(self.cs)

    def add_cs(self, w, cs):
        for server_param, client_param in zip(self.cs.parameters(), cs.parameters()):
            server_param.data += client_param.data.clone() * w

#条件选择模块，是一个1层神经网络
class ConditionalSelection(nn.Module):#条件选择模块
    def __init__(self, in_dim, h_dim):
        super(ConditionalSelection, self).__init__()
        #网络结构：1层线性层，归一化，ReLU
        self.fc = nn.Sequential(
            nn.Linear(in_dim, h_dim*2), 
            nn.LayerNorm([h_dim*2]),
            nn.ReLU(),
        )

    def forward(self, x, tau=1, hard=False):
        shape = x.shape
        x = self.fc(x)
        x = x.view(shape[0], 2, -1)
        x = F.gumbel_softmax(x, dim=1, tau=tau, hard=hard)
        return x[:, 0, :], x[:, 1, :]
